# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

from absl import app, flags

import datetime
import functools
import logging

from coin.proto.coin_feed_subscriber_pb2 import FeedSubscriberConfig
from coin.base.timestamp import get_timestamp
from coin.strategy.mm.base.unary_multifeed_strategy_base import UnaryMultifeedStrategyBase
from coin.strategy.mm.subscription import gen_feed_sub_request_with_prefer

from coin.strategy.mm.base.test_feed_arb_strategy import FeedInfo


class PrintFeedStrategy(UnaryMultifeedStrategyBase):
  def on_book(self, product, book):
    barray = book.get_bid_array(2)
    aarray = book.get_ask_array(2)
    bid1_price, bid1_qty = 0, 0
    ask1_price, ask1_qty = 0, 0
    if len(barray) > 1:
      bid1_price, bid1_qty = barray[1]
    if len(aarray) > 1:
      ask1_price, ask1_qty = aarray[1]

    print(
        FeedInfo(get_timestamp(),
                 'B',
                 book.timestamp,
                 product,
                 book.bid0().price,
                 book.bid0().qty,
                 book.ask0().price,
                 book.ask0().qty,
                 bid1_price,
                 bid1_qty,
                 ask1_price,
                 ask1_qty,
                 0,
                 0,
                 0))

  def on_trade(self, product, trade):
    print(
        FeedInfo(get_timestamp(),
                 'T',
                 trade.timestamp,
                 product,
                 0,
                 0,
                 0,
                 0,
                 0,
                 0,
                 0,
                 0,
                 trade.price,
                 trade.qty,
                 trade.side))

  def on_og_reset(self):
    pass

  def prepare(self):
    prod_to_sub = {}
    for key, values in self.feed_sub_req._sub_request_to_products_map.items():
      for value in values:
        prod_to_sub[value] = key

    for sub_req, products in self.feed_sub_req._sub_request_to_products_map.items():
      for product in products:
        feed_sub_request = gen_feed_sub_request_with_prefer(product, prod_to_sub[product])
        self.feed_subsystem.subscribe_book(feed_sub_request,
                                           functools.partial(self.on_book, product))
        self.feed_subsystem.subscribe_trade(feed_sub_request,
                                            functools.partial(self.on_trade, product))


class NullEnv(object):
  def __init__(self, *args, **kwargs):
    self.on_order_gateway_reset_callback = []

  def start(self):
    pass


def get_feed_strategy(feed_sub_req):
  ret_cls = PrintFeedStrategy

  products = feed_sub_req.all_products

  trade_product = products[0]
  ret_cls.CurrencyType = type(trade_product.base)
  ret_cls.ProductType = type(trade_product)
  ret_cls.LiveEnvironmentType = NullEnv

  ret = ret_cls(trade_product, [trade_product.base])
  ret.feed_sub_req = feed_sub_req
  return ret


def run_feed_strategy(*, feed_info, exit_after):
  import coin.strategy.mm.fastfeature.fast_feature_feed_spec as fffspec
  feedspec = fffspec.from_file(feed_info, datetime.datetime.now())
  strategy = get_feed_strategy(feedspec.feed_sub_request)

  strategy.run_live(strategy=strategy,
                    og_config_name=None,
                    on_order_gateway_reset_callback=strategy.on_og_reset,
                    exit_after=exit_after)


def main(_):
  assert flags.FLAGS.feed_info is not None
  run_feed_strategy(feed_info=flags.FLAGS.feed_info,
                    exit_after=datetime.timedelta(seconds=flags.FLAGS.exit_after_sec))


if __name__ == '__main__':
  flags.DEFINE_integer('num_feed_worker', None, '')
  flags.DEFINE_string('feed_info', None, '')
  flags.DEFINE_integer('exit_after_sec', 20, '')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
